
//
// 	NetFilterSDK
// 	Copyright (C) Vitaly Sidorov
//	All rights reserved.
//
//	This file is a part of the NetFilter SDK.
//	The code and information is provided "as-is" without
//	warranty of any kind, either expressed or implied.
//

#ifndef __TCP_PROXY_H__
#define __TCP_PROXY_H__
#include "stdafx.h"
#include <set>
#include <map>
#include <vector>
#include <mswsock.h>
#include <mstcpip.h>
#include "sync.h"
#include "socksdefs.h"
#include "iocp.h"
#include "threadpool.h"
#include "nfevents.h"
#include "proxydefs.h"

namespace TcpProxy
{

    class TCPProxy : public IOCPHandler, public ThreadJobSource
    {
    public:
        TCPProxy();
        virtual ~TCPProxy();
        unsigned short getPort();
        bool isIPFamilyAvailable(int ipFamily);
        bool init(unsigned short port, bool bindToLocalhost = true, int threadCount = 0);
        virtual void free();
        bool setProxy(NFAPI_NS ENDPOINT_ID id, PROXY_TYPE proxyType,
                      const char *proxyAddress, int proxyAddressLen,
                      const char *userName = NULL, const char *userPassword = NULL);
        void setEventHandler(NFAPI_NS NF_EventHandler *pEventHandler);
        void setTimeout(DWORD value);

        virtual bool getRemoteAddress(sockaddr *pRemoteAddr, NFAPI_NS PNF_TCP_CONN_INFO pConnInfo) = 0;
        bool tcpPostSend(NFAPI_NS ENDPOINT_ID id, const char *buf, int len);
        bool tcpPostReceive(NFAPI_NS ENDPOINT_ID id, const char *buf, int len);
        bool tcpSetState(NFAPI_NS ENDPOINT_ID id, bool suspended);
        bool tcpClose(NFAPI_NS ENDPOINT_ID id);
        bool getTCPConnInfo(NFAPI_NS ENDPOINT_ID id, NFAPI_NS PNF_TCP_CONN_INFO pConnInfo);

    protected:
        OV_DATA *newOV_DATA();
        void deleteOV_DATA(OV_DATA *pov);

        class AutoProxyData
        {
        public:
            AutoProxyData(TCPProxy *pThis, NFAPI_NS ENDPOINT_ID id)
            {
                thisClass = pThis;
                ptr = pThis->findProxyData(id);
                if (ptr)
                {
                    ptr->lock.Lock();
                }
            }
            ~AutoProxyData()
            {
                if (ptr)
                {
                    ptr->lock.Unlock();
                    thisClass->releaseProxyData(ptr);
                }
            }
            PROXY_DATA *operator->()
            {
                return ptr;
            }
            operator PROXY_DATA *()
            {
                return ptr;
            }
            TCPProxy *thisClass;
            PROXY_DATA *ptr;
        };
        PROXY_DATA *findProxyData(NFAPI_NS ENDPOINT_ID id);
        int releaseProxyData(PROXY_DATA *pd);
        void *getExtensionFunction(SOCKET s, const GUID *which_fn);
        bool initExtensions();
        bool startAccept(int ipFamily);
        bool startConnect(SOCKET socket, sockaddr *pAddr, int addrLen, NFAPI_NS ENDPOINT_ID id);
        bool startTcpSend(PROXY_DATA *pd, bool isInSocket, const char *buf, int len, NFAPI_NS ENDPOINT_ID id);
        bool startTcpReceive(PROXY_DATA *pd, bool isInSocket, NFAPI_NS ENDPOINT_ID id);
        bool startClose(SOCKET socket, NFAPI_NS ENDPOINT_ID id);
        void onConnectComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void onSendComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void onReceiveComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void setKeepAliveVals(SOCKET s);
        void onAcceptComplete(SOCKET socket, DWORD dwTransferred, OV_DATA *pov, int error);
        void onClose(SOCKET socket, OV_DATA *pov);
        virtual void onComplete(SOCKET socket, DWORD dwTransferred, OVERLAPPED *pOverlapped, int error);
        virtual void execute();
        virtual void threadStarted();
        virtual void threadStopped();

    private:
        IOCPService m_service;

        LIST_ENTRY m_ovDataList;
        int m_ovDataCounter;

        typedef std::map<NFAPI_NS ENDPOINT_ID, PROXY_DATA *> tSocketMap;
        tSocketMap m_socketMap;

        LPFN_ACCEPTEX m_pAcceptEx;
        LPFN_CONNECTEX m_pConnectEx;
        LPFN_GETACCEPTEXSOCKADDRS m_pGetAcceptExSockaddrs;

        SOCKET m_listenSocket;
        SOCKET m_acceptSocket;

        SOCKET m_listenSocket_IPv6;
        SOCKET m_acceptSocket_IPv6;

        NFAPI_NS ENDPOINT_ID m_connId;

        unsigned short m_port;

        bool m_ipv4Available;
        bool m_ipv6Available;

        NFAPI_NS NF_EventHandler *m_pPFEventHandler;
        LIST_ENTRY m_eventList;
        AutoCriticalSection m_csEventList;

        ThreadPool m_pool;

        DWORD m_timeout;

        PROXY_TYPE m_proxyType;
        char m_proxyAddress[NF_MAX_ADDRESS_LENGTH];
        int m_proxyAddressLen;

        std::string m_userName;
        std::string m_userPassword;

        AutoCriticalSection m_cs;
    };

    class LocalTCPProxy : public TCPProxy
    {
    public:
        void free();
        void setConnInfo(NFAPI_NS PNF_TCP_CONN_INFO pConnInfo);
        virtual bool getRemoteAddress(sockaddr *pRemoteAddr,
                                      NFAPI_NS PNF_TCP_CONN_INFO pConnInfo);

    private:
        typedef std::map<unsigned short, NFAPI_NS NF_TCP_CONN_INFO> tConnInfoMap;
        tConnInfoMap m_connInfoMap;

        AutoCriticalSection m_cs;
    };
}
#endif